import streamlit as st
import requests, time

st.title('文本分类预测')
text_input = st.text_area("请输入要预测的文本：", "中国人民公安大学2012年硕士研究生目录及书目")


def send_request(model_type):
    url = 'http://localhost:80/predict'
    data = {'text': text_input, 'model': model_type}
    start = time.time()
    try:
        resp = requests.post(url, json=data)
        cost_ms = (time.time() - start) * 1000
        if resp.status_code == 200:
            result = resp.json()['result']
            st.info(f'success cost={cost_ms:.2f}ms result={result}')
        else:
            st.error(f'error cost={cost_ms:.2f}ms status_code={resp.status_code}')
    except Exception as e:
        st.error(f'request error:{e}')


if st.button('predict by rf'):
    send_request('rf')

if st.button('predict by fasttext'):
    send_request('fasttext')

if st.button('predict by bert'):
    send_request('bert')

if st.button('predict by llm'):
    send_request('llm')
